from flask import Flask, request, jsonify
import pandas as pd
import os

from process.feature_process import *
from process.data_preprocess import *
from process.read_excel import read_data
from utils.model_utils import train_model, predict
from config import Config
from benchmark import find_best_method

app = Flask(__name__)

trained_model = None
feature_scaler = None
config = Config()

PREPROCESS_METHODS = {
    'pre_replace': pre_replace,
    'del_iqr': del_iqr,
    'del_zscore': del_zscore,
    'interpolate': interpolate,
    'mean': mean,
    'median': median
}

FEATURE_METHODS = {
    'ymsplit': ymsplit,
    'continuous_month': continuous_month,
    'seasonal': seasonal,
    'sliding_window': sliding_window,
    'combined': combined,
    'seasonal_plus': seasonal_plus
}

def init_model(force_update=False, preprocess_method=None, feature_method=None):
    """初始化并训练模型
    
    Args:
        force_update: 是否强制更新模型
        preprocess_method: 指定的预处理方法名
        feature_method: 指定的特征工程方法名
    """
    global trained_model, feature_scaler
    
    df = read_data(os.path.join(os.path.dirname(__file__), config.config['data_path']), is_new_format=True)

    if preprocess_method is not None:
        if preprocess_method not in PREPROCESS_METHODS:
            raise ValueError(f"无效的预处理方法: {preprocess_method}。可用方法: {list(PREPROCESS_METHODS.keys())}")
    if feature_method is not None:
        if feature_method not in FEATURE_METHODS:
            raise ValueError(f"无效的特征工程方法: {feature_method}。可用方法: {list(FEATURE_METHODS.keys())}")
    
    # 如果没有指定方法且需要更新，则自动选择最佳方法
    if (preprocess_method is None or feature_method is None) and (force_update or config.is_data_modified()):
        print("自动选择最佳方法...")
        auto_preprocess, auto_feature, best_score = find_best_method(df)
        preprocess_method = preprocess_method or auto_preprocess
        feature_method = feature_method or auto_feature
        print(f"选择结果 - 预处理: {preprocess_method}, 特征工程: {feature_method}, 得分: {best_score:.4f}")
        config.update_best_method(preprocess_method, feature_method)
    
    # 使用选定的方法
    preprocess_module = PREPROCESS_METHODS[preprocess_method]
    feature_module = FEATURE_METHODS[feature_method]
    # 更新最后修改时间
    config.update_last_modified()
    
    # 处理并训练
    df = preprocess_module.process(df)
    if df is None:
        raise Exception("数据预处理失败")
    X, y = feature_module.prepare_features(df)
    model, _, _, _, _ = train_model(X, y, test_size=0.1, random_state=114514)
    trained_model = model
    feature_scaler = feature_module

@app.route('/api/predict', methods=['POST'])
def predict_usage():
    """用电量预测API"""
    try:
        data = request.get_json()
        if not data or 'date' not in data:
            return jsonify({'error': '缺少date参数'}), 400
        
        preprocess_method = data.get('preprocess')
        feature_method = data.get('feature')
        
        force_update = False
        # 如果指定了，使用指定的
        if preprocess_method or feature_method:
            force_update = True
            # 更新last_method
            config.update_last_method(
                preprocess_method or config.config['best_method']['preprocess'],
                feature_method or config.config['best_method']['feature']
            )
        # 如果没有指定，使用best_method
        else:
            # 上次使用的不同，需要重新初始化
            if (config.config['last_method']['preprocess'] != config.config['best_method']['preprocess'] or
                config.config['last_method']['feature'] != config.config['best_method']['feature']):
                force_update = True
            # 使用best_method并更新last_method
            preprocess_method = config.config['best_method']['preprocess']
            feature_method = config.config['best_method']['feature']
            config.update_last_method(preprocess_method, feature_method)
        
        # 需要重新初始化模型
        if force_update or config.is_data_modified():
            print(f"重新初始化模型 - 预处理: {preprocess_method}, 特征工程: {feature_method}")
            init_model(force_update=True, 
                      preprocess_method=preprocess_method,
                      feature_method=feature_method)
            
        # 创建预测用的数据框
        target_date = pd.to_datetime(data['date'])
        predict_df = pd.DataFrame({
            0: [target_date.strftime('%Y-%m')],
            1: [0]
        })
        # 特征
        X_pred = feature_scaler.prepare_features(predict_df, is_predict=True)
        prediction = predict(trained_model, X_pred)
        
        return jsonify({
            'date': data['date'],
            'predicted_usage': float(prediction[0]),
            'current_method': {
                'preprocess': config.config['last_method']['preprocess'],
                'feature': config.config['last_method']['feature']
            },
            'last_data_modified': config.config['last_modified']
        })    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # 启动时初始化模型
    init_model(force_update=True)
    app.run(host='0.0.0.0', port=5000) 